import uuid
import json
import os
from datetime import datetime
from pathlib import Path
from typing import Dict, Optional, List
import logging
from config import BASE_DIR
from utils.gpu_monitor import GPUMonitor

class TaskManager:
    def __init__(self):
        self.tasks_dir = BASE_DIR / "tasks"
        self.logs_dir = BASE_DIR / "logs"
        self.data_dir = BASE_DIR / "data" / "custom"
        
        # 创建必要的目录
        for dir_path in [self.tasks_dir, self.logs_dir, self.data_dir]:
            dir_path.mkdir(parents=True, exist_ok=True)
        
        # 配置日志
        logging.basicConfig(level=logging.INFO)
        self.logger = logging.getLogger(__name__)
        
        # 加载现有任务
        self.tasks: Dict[str, dict] = self._load_tasks()
        
        # GPU监控器
        self.gpu_monitor = GPUMonitor()
    
    def _load_tasks(self) -> Dict[str, dict]:
        tasks = {}
        tasks_file = self.tasks_dir / "tasks.json"
        if tasks_file.exists():
            with open(tasks_file, "r") as f:
                tasks = json.load(f)
        return tasks
    
    def _save_tasks(self):
        with open(self.tasks_dir / "tasks.json", "w") as f:
            json.dump(self.tasks, f, indent=2)
    
    def create_task(self, task_type: str, params: dict) -> str:
        """创建新任务并返回任务ID"""
        # 设置默认参数
        default_params = {
            "epochs": 50,
            "learning_rate": 0.001,
            "batch_size": 32,
            "optimizer": "adam",
            "dataset_path": None
        }
        
        if task_type == "time_series":
            default_params.update({
                "sequence_length": 96,
                "prediction_length": 24
            })
        elif task_type == "image_classification":
            default_params.update({
                "num_classes": 2
            })
        
        # 更新参数
        params = {**default_params, **params}
        
        task_id = str(uuid.uuid4())
        task_info = {
            "id": task_id,
            "type": task_type,
            "params": params,
            "status": "created",
            "created_at": datetime.now().isoformat(),
            "updated_at": datetime.now().isoformat(),
            "metrics": [],
            "gpu_status": self.gpu_monitor.get_gpu_status()
        }
        
        self.tasks[task_id] = task_info
        self._save_tasks()
        
        # 创建任务日志文件
        log_file = self.logs_dir / f"{task_id}.log"
        handler = logging.FileHandler(log_file)
        handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
        
        task_logger = logging.getLogger(task_id)
        task_logger.addHandler(handler)
        task_logger.setLevel(logging.INFO)
        
        return task_id
    
    def update_task_status(self, task_id: str, status: str, metrics: Optional[dict] = None):
        """更新任务状态和指标"""
        if task_id not in self.tasks:
            raise ValueError(f"Task {task_id} not found")
        
        self.tasks[task_id]["status"] = status
        self.tasks[task_id]["updated_at"] = datetime.now().isoformat()
        self.tasks[task_id]["gpu_status"] = self.gpu_monitor.get_gpu_status()
        
        if metrics:
            self.tasks[task_id]["metrics"].append({
                "timestamp": datetime.now().isoformat(),
                **metrics
            })
        
        self._save_tasks()
        
        # 记录到日志
        task_logger = logging.getLogger(task_id)
        log_message = f"Status: {status}"
        if metrics:
            log_message += f" | Metrics: {metrics}"
        task_logger.info(log_message)
    
    def get_task_info(self, task_id: str) -> dict:
        """获取任务信息"""
        if task_id not in self.tasks:
            raise ValueError(f"Task {task_id} not found")
        
        # 更新GPU状态
        self.tasks[task_id]["gpu_status"] = self.gpu_monitor.get_gpu_status()
        return self.tasks[task_id]
    
    def get_task_log(self, task_id: str) -> str:
        """获取任务日志"""
        log_file = self.logs_dir / f"{task_id}.log"
        if not log_file.exists():
            raise ValueError(f"Log file for task {task_id} not found")
        
        with open(log_file, "r") as f:
            return f.read()
    
    def save_dataset(self, file_content: bytes, dataset_type: str, dataset_name: str) -> Dict:
        """保存上传的数据集"""
        dataset_id = f"{dataset_type}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        dataset_dir = self.data_dir / dataset_id
        dataset_dir.mkdir(exist_ok=True)
        
        file_path = dataset_dir / dataset_name
        with open(file_path, "wb") as f:
            f.write(file_content)
        
        return {
            "dataset_id": dataset_id,
            "path": str(file_path),
            "type": dataset_type
        } 